import argparse, math

protein_polarity = {
    'A': 8.1,  #/* A */
	'R': 10.5,  #/* R */
	'N': 11.6,  #/* N */
	'D': 13.0,  #/* D */
	'C': 5.5,  #/* C */
	'Q': 10.5,  #/* Q */
	'E': 12.3,  #/* E */
	'G': 9.0,  #/* G */
	'H': 10.4,  #/* H */
	'I': 5.2,  #/* I */
	'L': 4.9,  #/* L */
	'K': 11.3,  #/* K */
	'M': 5.7,  #/* M */
	'F': 5.2,  #/* F */
	'P': 8.0,  #/* P */
	'S': 9.2,  #/* S */
	'T': 8.6,  #/* T */
	'W': 5.4,  #/* W */
	'Y': 6.2,  #/* Y */
	'V': 5.9  #/* V */
}

protein_volume = {
	'A': 31.0,  #/* A */
	'R': 124.0,  #/* R */
	'N': 56.0,  #/* N */
	'D': 54.0,  #/* D */
	'C': 55.0,  #/* C */
	'Q': 85.0,  #/* Q */
	'E': 83.0,  #/* E */
	'G':  3.0,  #/* G */
	'H': 96.0,  #/* H */
	'I': 111.0,  #/* I */
	'L': 111.0,  #/* L */
	'K': 119.0,  #/* K */
	'M': 105.0,  #/* M */
	'F': 132.0,  #/* F */
	'P': 32.5,  #/* P */
	'S': 32.0,  #/* S */
	'T': 61.0,  #/* T */
	'W': 170.0,  #/* W */
	'Y': 136.0,  #/* Y */
	'V': 84.0  #/* V */
}

pi = math.acos(-1.0)

def parse_fasta(file_buffer):
    '''
    ignore the description line and return the sequence line, the file must have only 2 lines\n
    returns: string
    exception: TypeError: Check the file is illegal
    '''
    cnt = 0
    for data in file_buffer:
        if cnt == 1: return data
        cnt += 1
    # If the file only have 1 line, then raise an Error
    raise TypeError("File is illegal!")

def argparser():
    '''
    define a parser to process file name\n
    return: processed arguments
    '''
    parser = argparse.ArgumentParser(description = "Recurrent the FFT processing module in MAFFT. The file format is fasta.")
    parser.add_argument("--file1", type = str, help = "Input file 1")
    parser.add_argument("--file2", type = str, help = "Input file 2")
    parser.add_argument("-n", "--normalize-protein", type = int, choices = [0, 1], default = 1, help = "Normalize protein data, 1 is True, 0 is False")
    return parser.parse_args()
    
def protein_to_polarity(s: str):
    '''
    given a protein sequence, calcuate the polarity sequence
    return: list with protein polarity values sequence
    '''
    L = []
    for i in s:
        L.append(protein_polarity[i])
    return L

def protein_to_volume(s: str):
    '''
    given a protein sequence, calcuate the volume sequence
    return: list with protein volume values sequence
    '''
    L = []
    for i in s:
        L.append(protein_volume[i])
    return L

rev = []
coss = []
sinn = []

def init_ave_std() -> None:
    '''
    Calcuate the Average and Standard of the protein polarity and volume
    return: None
    '''
    global ave_polar, ave_volume, std_polar, std_volume
    ave_polar = ave_volume = std_polar = std_volume = 0.0
    for i in protein_polarity:
        ave_polar += protein_polarity[i]
    ave_polar /= len(protein_polarity)
    for i in protein_polarity:
        std_polar += (protein_polarity[i] - ave_polar) ** 2
    std_polar /= len(protein_polarity)
    std_polar = math.sqrt(std_polar)
    for i in protein_polarity: protein_polarity[i] = (protein_polarity[i] - ave_polar) / std_polar
    # print(ave_polar, std_polar, protein_polarity)
    for i in protein_volume:
        ave_volume += protein_volume[i]
    ave_volume /= len(protein_volume)
    for i in protein_volume:
        std_volume += (protein_volume[i] - ave_volume) ** 2
    std_volume /= len(protein_volume)
    std_volume = math.sqrt(std_volume)
    for i in protein_volume: protein_volume[i] = (protein_volume[i] - ave_volume) / std_volume
    # print(ave_volume, std_volume, protein_volume)
    return 
        

def init_seq() -> None:
    '''
    Reset rev, coss, sinn array
    return: None
    '''
    rev = []
    coss = []
    sinn = []

def fft(length: int, a: complex, o: int):
    '''
    FFT Function, make the array A change the place
    return: Transformed array
    '''
    for i in range(length):
        if i < rev[i]: 
            a[i], a[rev[i]] = a[rev[i]], a[i]
    j = 1
    while j < length:
        wn = complex(math.cos(pi / j), math.sin(pi / j) * o)
        for k in range(0, length, j << 1):
            w0 = complex(1.0, 0.0)
            for i in range(0, j):
                X = a[i + k]
                Y = w0 * a[i + j + k]
                a[i + k] = X + Y
                a[i + k + j] = X - Y
                w0 = w0 * wn
            if k >= length: break 
        j <<= 1
    if o == 1: 
        a = [i / length for i in a]
    return a


def converter(seq: str, Type: str, dis: int = 0):
    '''
    Convert Protein alphabet to float number.
    return: A List with float data which has polarity or volume data.
    Exception: TypeError: The type must in ['polarity', 'volume'].
    '''
    if Type == 'polarity': L = protein_to_polarity(seq[dis : len(seq)])
    elif Type == 'volume': L = protein_to_volume(seq[dis : len(seq)])
    else: raise TypeError("The type must in ['polarity', 'volume'].")
    return L

def complex_conjugate_product(a: complex, b: complex) -> complex:
    '''
    Calcuate the conjugate product. Let a be (x + jy), b be (w + jv), the product is (xw + yv) + j(wy - xv)
    return: product
    '''
    #print(a, b, complex(a.real * b.real + a.imag * b.imag, a.imag * b.real - a.real * b.imag))
    return complex(a.real * b.real + a.imag * b.imag, a.imag * b.real - a.real * b.imag)

def fft_main(seq1: str, seq2: str):
    '''
    FFT Main, process the sequence to float num and calc \sum_{1 <= i <= n, 1 <= j <= n + m} {seq1(i) * seq2(i + k)} 
    return: Calcuation
    Exception: 
    '''
    global length, len1, len2
    l11 = converter(seq1, 'polarity')
    l12 = converter(seq2, 'polarity')
    l21 = converter(seq1, 'volume')
    l22 = converter(seq2, 'volume')
    len1 = len(l11)
    len2 = len(l12)
    if len1 <= 0 or len2 <= 0: raise TypeError("The length of sequence is less than 0")
    length = 1
    l = 0
    while(length <= max(len1, len2)):
        length <<= 1
        l += 1
    # Add 0
    for i in range(len1, length + 1): l11.append(0); l21.append(0)
    for i in range(len2, length + 1): l12.append(0); l22.append(0)
    # Change to imag data
    cpx1 = []
    cpx2 = []
    for kssn in range(length): cpx1.append(complex(l11[kssn], l21[kssn]))
    for kssn in range(length): cpx2.append(complex(l12[kssn], l22[kssn]))
    rev.append(0)
    for i in range(1, length + 1):
        rev.append(((rev[i >> 1] >> 1) | ((i & 1) << (l - 1))))
    ##### Polygon multiply polygon
    #print(cpx2)
    cpx1 = fft(length, cpx1, 1) # DFT
    cpx2 = fft(length, cpx2, 1)
    #print(cpx2)
    ans = []
    for i in range(length): ans.append(complex_conjugate_product(cpx1[i], cpx2[i]))
    #print(ans)
    fft(length, ans, -1) # IDFT
    #print(ans)
    return ans

if __name__ == '__main__':
    args = argparser()
    f1 = open(args.file1, "r")
    f2 = open(args.file2, "r")
    s1 = parse_fasta(f1)
    s2 = parse_fasta(f2)
    if args.normalize_protein: init_ave_std()

    protein_ans = fft_main(s1, s2)
    print("Protein property: ")
    cnt = 0
    for i in protein_ans:
        cnt += 1
        print('c(%d) = %.5f' % (cnt - len2, i.real))
    